/**
  ETFAna project, Anyang Normal University && IMP-CAS
  \class ETFPIDm
  \brief inherited from ETFVPID - 
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2023-06-02
  \date 2023-06-02 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2023-06-02   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021-2024 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <TGraph.h>
#include <TLegend.h>
#include <TRandom3.h>
#include <TClonesArray.h>
#include <TTree.h>
#include "ETFRun.h"
#include "ETFPIDm.h"
#include "ETFMsg.h"
#include "ETFDetectorManager.h"
#include "ETFTrackManager.h"
#include "ETFPostMagTracker.h"
#include "ETFMath.h"
#include "ETFTOFWall.h"
#include "ETFDetector.h"
#include "ETFTrack.h"
#include "ETF3DTrack.h"
#include "ETFTimer2.h"
#include "ETFPlaHit.h"
#include "ETFParticle.h"
#include "ETFParaManager.h"

using std::cout;
using std::endl;

#define unisolv ETFMath::UniformMagneticSolution
#define sqr ETFMath::sqr

ETFPIDm::ETFPIDm() : ETFVPID(){
  fParticle = new ETFParticle();

  fpidv = new TClonesArray("ETFParticle");
  fpidv->SetName("pidm");
} // end ctor

ETFPIDm::~ETFPIDm(){
  if(fpidv) delete fpidv;
  if(fParticle) delete fParticle;
} // end dtor

const ETFParticle *ETFPIDm::GetParticle(int i) const{
  if(i >= int(fpidv->GetEntriesFast()))
    ETFMsg::Error("ETFPIDm", "GetParticle: requested i=%d out of bounds", i);
  
  ETFParticle *p = dynamic_cast<ETFParticle *>(fpidv->At(i));
  if(!p) ETFMsg::Error("ETFPIDm", "GetParticle: the requird paticle is nullptr");
  
  return p;
} // end member function GetParticle

int ETFPIDm::GetNofParticles() const{
  return fpidv->GetEntriesFast();
} // end member function GetNofParticles

void ETFPIDm::Initialize(){
  ETFVPID::Initialize();
  fTaHitX = -999.;
  fParticle->Initialize();
  fpidv->Clear();
} // end member function Initialize

// register to the main tree
void ETFPIDm::Register(){
  fTree->Branch(fpidv->GetName(), &fpidv);
} // end member function Register

// get the hit position on the target, i.e., assigning fTaHitX
// return true on success
bool ETFPIDm::GetTaHitPos(){
  // for sipmArr, 2 situations are allowed:
  // 1. only 1 strip fired; 2. 2 adjacent strips fired //
  static ETFTOFWall *sipmarr = dynamic_cast<ETFTOFWall *>(fdm->GetDetector("sipmArr", true));
  const int ns = sipmarr->GetNFiredStripSingle(); // multiplicity of SiPMArr
  if(!ns || ns > 2) return false; // no SiPMArr strips are fired, or too many are fired
  
  // obtain the target hit position //
  auto &mm = *sipmarr->hits();
  fTaHitX = 0.;
  int stripId = -1;
  for(TObject *b : mm){ // filter those spurious hits
    auto h = static_cast<ETFPlaHit *>(b);
    if(!h || !h->GetFiredStatus()) continue; // only the TOT channel fired -- a spurious hit
    if(1 == ns){ fTaHitX = h->GetX(); return true; }
    if(stripId != -1 && abs(stripId - h->stripId()) != 1){ // not adjacent strip, so drop it
      fTaHitX = -999.;
      return false;
    } // end if
    stripId = h->stripId();
    fTaHitX += h->GetX();
  } // end for over hits

  fTaHitX /= 2.; // only two hit strip can reach here
	// FIXME BINGER:TASiPMPlaArray.cxx:172  ?? FIXME
  // ??? fX0 += 3.; // 0 strip is broken, so 0 strip now is strip 1, 20171027_1445 ???
  return true;
} // end member function GetTaHitPos


bool ETFPIDm::Fly(bool istof){
  if(fHasFlied) return true;

  static ETFPostMagTracker *pmt = ftm->GetPostMagTracker();
  static ETFPostMagTracker *mpmt = ftm->GetmPostMagTracker();
  if(!GetTaHitPos()) return false; // get target hit position

  // now we can do the PID //
  static bool is3D = fpm->is3DTrk(); // whether using 3D or not
  if(pmt) fFlySuccess = Fly(pmt, is3D, istof); // the main tracker
  if(mpmt) fFlySuccess = fFlySuccess || Fly(mpmt, is3D, istof); // the minor tracker

  return fFlySuccess;
} // end member function Fly

// the general fly function that will actually be called by the program
// istof: whether we only call fly to calculate tof
// is3D: whether using 3D tracks
// trker specifies which tracker to select tracks from
bool ETFPIDm::Fly(ETFPostMagTracker *trker, bool is3D, bool istof){
  static ETFTimer2 *tstop = dynamic_cast<ETFTimer2 *>(fdm->GetDetector("tstop", true));
  if(4 != tstop->GetFiredStatus()) return false;
  ETFTOFWall *tofw = trker->GetTOFWall(); // DO NOT USE "static" HERE -- changes with trker

  // get the postmagtrk of the light fragments //
  // NOTE that there may be multiple tracks in one event which entails a loop
  int ntr = 0;
  if(is3D) ntr = trker->GetNof3DTracks();
  else ntr = trker->GetNofTracks();
  if(!ntr) return false;

  // loop over each postmagtrk for PID //
  for(int i = 0; i < ntr; i++){
    fParticle->dcArrId = bool(strcmp(trker->GetName(), "mpostMagTracker"));
    // calculate tof2 //
    ETFTrack *trkX = nullptr; ETF3DTrack *trk3d = nullptr;
    if(is3D){
      trk3d = trker->Get3DTrack(i);
      trkX = dynamic_cast<ETFTrack *>(trk3d->GetTrack(TrackType::kX));
    } // end if
    else trkX = dynamic_cast<ETFTrack *>(trker->GetTrack(i, TrackType::kX));

    const int stripId = trkX->GetFiredStripId();
    fParticle->tof2 = tofw->GetStrip(stripId)->GetTime() - tstop->GetTime() + ftof2_ofst;
    if(istof) return false;

    fParticle->k1 = trkX->GetK(); fParticle->b1 = trkX->GetB();
    if(is3D){ fParticle->k2 = trk3d->GetK2(); fParticle->b2 = trk3d->GetB2(); }
    if(Fly(tofw)){
      ETFParticle *p = dynamic_cast<ETFParticle *>(fpidv->ConstructedAt(fpidv->GetEntriesFast()));
      *p = *fParticle;
    } // end if
    else{
      fParticle->Initialize();
      return false;
    } // end else
  } // end loop over postmagtrks

  fHasFlied = true;
  return true;
} // end member function Fly

// propogate the particle in the dipole magnet from the target-hit point
// to meet the post-mag trk
bool ETFPIDm::Fly(const ETFTOFWall *tofw){  
  // solve the particle radius in the magfield using point+trk
  // dtheta: deflection angle; rho, (zo, xo): defines the arc;
  // (ki, bi): defines the incident line; (x1, x2): x's of points of into and from magField
  double result[8]{}; // [0-7]: dtheta, rho, (ki, bi), (zo, xo), (x1, x2)
  if(!unisolv(fParticle->k1, fParticle->b1, kzMagOut, kzMagIn, fTaZ, fTaHitX, result))
    return false;
  fParticle->dtheta = result[0];
  if(-999. == kB) ETFMsg::Error("ETFPID", "Fly: kB not assigned");
  fParticle->rho = result[1];
  fParticle->brho = fParticle->rho*kB*1.e-3; // in T.m; r: 1e-3 factor for mm->m
  // fParticle->brho = 0.760226*fParticle->brho + 0.012984; // calibratoin - Jun. 15, 2023 //
  fParticle->ki = result[2]; fParticle->bi = result[3];
  fParticle->zo = result[4]; fParticle->xo = result[5];
  fParticle->x1 = result[6]; fParticle->x2 = result[7];
  // calculate the total flight path corresponding to tof2 //
  const double pi[4] = {fParticle->k1, fParticle->b1, 0., 0.}; // [0-4]: {k, b, k2, b2}
  double pHit[3]{}; // store the hit point [0-1-2]: [x-y-z]
  tofw->GetHitPosition(pi, pHit);
  // total track length = Ta->magIn + arc + magOut->TOFWall
  fParticle->trkLenT = (kzMagIn-fTaZ)*sqrt(1.+sqr(fParticle->ki)) +
    fParticle->rho*fParticle->dtheta + (pHit[2]-kzMagOut)*sqrt(1.+sqr(fParticle->k1)); // in mm
  if(-999. != fParticle->k2) fParticle->trkLenT *= sqrt(1.+sqr(fParticle->k2));
  fParticle->trkLenT += (fTaZ-fTstopZ); // += tstop->Ta
  const double b = fParticle->trkLenT / fParticle->tof2 / c0;
  if(b < 0. || b > 1.) return false;
  fParticle->beta = b; fParticle->bm2 = pow(b, -2.);
  fParticle->z = 1.; // FIXME:FIXME --  temporarily set to 1  -- FIXME:FIXME
  const double bg = ETFMath::BetaGamma(b);
  fParticle->gamma = bg/b;
  if(-999. == kB) ETFMsg::Error("ETFPID", "Fly: kB not assigned");
  // CC = 0.321840605 = e0/(u0*c0*1E6) SI unit
  fParticle->aoz = fParticle->brho*CC/bg;
  fParticle->dtheta /= ETFMath::DEGREE();
  // calculate the projection of poz w.r.t. incident direction: n(nx,ny,nz)=(k1, k2, 1) //
  fParticle->poz = fParticle->brho*c0; // in MeV/c
  double tmp = sqr(fParticle->k1);
  if(-999. != fParticle->k2) tmp += sqr(fParticle->k2);
  fParticle->pozp = fParticle->poz / sqrt(1. + tmp);
  fParticle->pozt = fParticle->pozp * sqrt(tmp);
  fParticle->index = fdm->Index();

  return true;
} // end member function Fly

void ETFPIDm::ConstructGraph(){
  ETFVPID::ConstructGraph();

  static const int NPT = 600; // nof points to draw
  static TRandom3 rdm; ///< for generating random points
 
  if(!fHasFlied || !GetNofParticles()) return; // no eligible for drawing
  if(!fGCurve) ETFMsg::Error("ETFPIDm", "ConstructGraph: fGCurve is nullptr");
  if(!fGLegend) ETFMsg::Error("ETFPIDm", "ConstructGraph: fGLegend is nullptr");

  char name[64]; // for display the aoz
  for(const auto &b : *fpidv){
    const int np = fGCurve->GetN(); // # points already in fGCurve
    int n = np;
    ETFParticle *p = dynamic_cast<ETFParticle *>(b);
    while(n - np <= NPT){
      const double ai = rdm.Uniform(-PI, PI); // (-pi, pi)
      const double zi = p->zo + p->rho*cos(ai);
      const double xi = p->xo + p->rho*sin(ai);
      // if(zi < kzMagIn || zi > kzMagOut) continue; // only draw the arc
      fGCurve->SetPoint(n++, zi, xi);
    } // end for over i
    // label the aoz value in the graph
    sprintf(name, "aoz_%1.3f\n", p->aoz);
    fGLegend->AddEntry(fGCurve, name, "lp");
  } // end loop over particles
} // end member function ConstructGraph


void ETFPIDm::Print(){
  cout << "\033[32;1m============= PIDm PRINT =============\033[0m" << endl;
  const int np = GetNofParticles();
  if(!np || -999. == fTaHitX){
    cout << "No particle is identified." << endl;
    return;
  } // end if
  cout << "target hit position: " << fTaHitX << " mm" << endl;
  cout << "number of particles " << np << endl;
  int i = 0;
  for(const auto &p : *fpidv){
    cout << "\033[32;1m\t----------- PARTICLE #" << i++ << " ---------------\033[0m" << endl;
    dynamic_cast<ETFParticle *>(p)->Print();
  } // end if
} // end member function Print
